def timeit(f):
    """Take a function f and return a new function that calls 5000 times,
    returning the last result and the average runtime, ignoring the ten 
    slowest/fastest calls. This gives a more reliable estimate for fast functions.
    """
    _time=time.time
    def timed_f(*args, **kwargs):
        times = []
        for _ in range(5000):
            start = _time()
            result = f(*args, **kwargs)
            end = _time()
            times.append(end - start)
        # Take the average of the middle three to smooth out variance from cache effects.
        average_time = sum(sorted(times)[10:-10]) / (len(times) - 20)
        return result, average_time
    return timed_f

def beta_v0(spy, asset):
    """Calculate the slope of a linear regression for one asset with SPY.
    """
    asset_residuals = asset - asset.mean()
    spy_residuals = spy - spy.mean()
    
    covariance = (asset_residuals * spy_residuals).sum() / (len(asset) - 1)
    spy_variance = (spy_residuals ** 2).sum() / (len(asset) - 1)
    
    return covariance / spy_variance 
    
v0_result, v0_duration = timeit(beta_v0)(rets[SPY].values, rets[AAPL].values)
print("Our AAPL beta is {}. Calculation took {} seconds.".format(v0_result, v0_duration))

scipy_result, scipy_duration = timeit(linregress)(rets[SPY].values, rets[AAPL].values)
print("linregress AAPL beta is {.slope}. Calculation took {} seconds.".format(
    scipy_result, scipy_duration
))
print()
print("Our calculation was {} times faster than linregress.".format(scipy_duration / v0_duration))

def beta_v1(spy, asset):
    """Calculate the slope of a linear regression for one asset with SPY.
    """
    asset_residuals = asset - asset.mean()
    spy_residuals = spy - spy.mean()
    
    # These variable names aren't really correct anymore since we're
    # taking a sum instead of a mean, but I couldn't think of better names.
    covariance = (asset_residuals * spy_residuals).sum()  # Not dividing here anymore.
    spy_variance = (spy_residuals ** 2).sum()             # Not dividing here anymore.
    
    return covariance / spy_variance

v1_result, v1_duration = timeit(beta_v1)(rets[SPY].values, rets[AAPL].values)
print("Our AAPL beta is still {}. Calculation took {} seconds.".format(v1_result, v1_duration))
print("Our calculation was {} times faster than linregress.".format(scipy_duration / v1_duration))

def vectorized_beta(spy, assets):
    """Calculate beta between every column of ``assets`` and ``spy``.
    
    Parameters
    ----------
    spy : np.array
        An (n x 1) array of returns for SPY.
    assets : np.array
        An (n x m) array of returns for m assets.
    """
    assert len(spy.shape) == 2 and spy.shape[1] == 1, "Expected a column vector for spy."

    asset_residuals = assets - assets.mean(axis=0)
    spy_residuals = spy - spy.mean()

    covariances = (asset_residuals * spy_residuals).sum(axis=0)
    spy_variance = (spy_residuals ** 2).sum()
    return covariances / spy_variance

from quantopian.pipeline.experimental import QTradableStocksUS

def get_qtradable_for(day):
    """Get the QTradableStocksUS for a single day.
    """
    pipe = Pipeline({}, screen=QTradableStocksUS())
    result = run_pipeline(pipe, day, day)
    return result.index.get_level_values(1)

universe = get_qtradable_for(END)
print(universe[:5])
Index([ Equity(2 [ARNC]), Equity(24 [AAPL]), Equity(41 [ARCB]),
        Equity(52 [ABM]),  Equity(62 [ABT])],
      dtype='object')

all_returns = returns(universe, START, END).dropna(how='any', axis=1)
all_returns.head()

# Convert 1D array of spy returns into an N x 1 2D array.
SPY_column = rets[SPY].values[:, np.newaxis]
betas, duration = timed(vectorized_beta)(SPY_column, all_returns.values)

print("AAPL beta is still", betas[all_returns.columns.get_loc(AAPL)])
print("It took {} seconds to calculate betas for {} stocks.".format(duration, len(betas)))

0.002 * 252

SPY_asset = symbols('SPY')

daily_returns = DailyReturns()

class MyBeta(CustomFactor):
    # Get daily returns for every asset in existence, plus the daily returns for just SPY
    # as a column vector.
    inputs = [daily_returns, daily_returns[SPY_asset]]
    # Set a default window length of 2 years.
    window_length = 252
    
    def compute(self, today, assets, out, all_returns, spy_returns):
        out[:] = vectorized_beta(spy_returns, all_returns)

my_beta_pipeline = Pipeline({'beta': MyBeta()})
my_beta_pipeline.show_graph('png')

my_betas, duration = timed(run_pipeline)(my_beta_pipeline, END, END)
print("AAPL's beta is still", my_betas.xs(AAPL, level=1).values[0, 0])
print("It took {} seconds to calculate 252-day betas for 1 day:".format(duration))

from zipline.utils.calendars import get_calendar
NYSE = get_calendar('NYSE')

my_beta_result, duration = timed(run_pipeline)(my_beta_pipeline, END, END + (5 * NYSE.day))
print("It took {} seconds to calculate 252-day betas for 5 trading days:".format(duration))

my_beta_result, duration = timed(run_pipeline)(my_beta_pipeline, END, END + (40 * NYSE.day))
print("It took {} seconds to calculate 252-day betas for two months:".format(duration))

my_beta_result, duration = timed(run_pipeline)(my_beta_pipeline, END, END + (252 * NYSE.day))
print("It took {} seconds to calculate 252-day betas for a year:".format(duration))

from quantopian.pipeline.factors import SimpleBeta
print(SimpleBeta.__doc__)

Factor producing the slope of a regression line between each asset's daily
    returns to the daily returns of a single "target" asset.

    Parameters
    ----------
    target : zipline.Asset
        Asset against which other assets should be regressed.
    regression_length : int
        Number of days of daily returns to use for the regression.
    allowed_missing_percentage : float, optional
        Percentage of returns observations that are allowed to be missing when
        calculating betas. Assets with more than this percentage of returns
        observations missing will produce values of NaN. Default behavior is
        that 25% of inputs can be missing.

SPY = symbols('SPY')
simple_beta = SimpleBeta(SPY, 252, allowed_missing_percentage=0.3)
simple_beta_pipeline = Pipeline({'beta': simple_beta})

simple_beta_result, duration = timed(run_pipeline)(simple_beta_pipeline, END, END + (252 * NYSE.day))
print("It took {} seconds to calculate 252-day betas for a year:".format(duration))

print("SimpleBeta NaN count for 2010:", simple_beta_result.beta.isnull().sum())
print("MyBeta NaN count for 2010:", my_beta_result.beta.isnull().sum())